#' Perform 2 sample MR on each SNP individually
#'
#' @param dat Output from [harmonise_data()].
#' @param parameters List of parameters. The default is `default_parameters()`.
#' @param single_method Function to use for MR analysis. The default is `"mr_wald_ratio"`.
#' @param all_method Functions to use for MR analysis. The default is `c("mr_ivw", "mr_egger_regression")`.
#'
#' @export
#' @return List of data frames
mr_singlesnp <- function(
  dat,
  parameters = default_parameters(),
  single_method = "mr_wald_ratio",
  all_method = c("mr_ivw", "mr_egger_regression")
) {
  if (!"samplesize.outcome" %in% names(dat)) {
    dat$samplesize.outcome <- NA
  }

  stopifnot("outcome" %in% names(dat))
  stopifnot("exposure" %in% names(dat))
  stopifnot("beta.exposure" %in% names(dat))
  stopifnot("beta.outcome" %in% names(dat))
  stopifnot("se.exposure" %in% names(dat))
  stopifnot("se.outcome" %in% names(dat))

  res <- plyr::ddply(dat, c("id.exposure", "id.outcome"), function(X) {
    x <- subset(X, mr_keep)
    nsnp <- nrow(x)
    if (nsnp == 0) {
      x <- X[1, ]
      d <- data.frame(
        SNP = "No available data",
        b = NA,
        se = NA,
        p = NA,
        samplesize = NA,
        outcome = x$outcome[1],
        exposure = x$exposure[1]
      )
      return(d)
    }
    l <- lapply(1:nsnp, function(i) {
      with(
        x,
        get(single_method)(
          beta.exposure[i],
          beta.outcome[i],
          se.exposure[i],
          se.outcome[i],
          parameters
        )
      )
    })
    nom <- c()
    for (i in seq_along(all_method)) {
      l[[nsnp + i]] <- with(
        x,
        get(all_method[i])(beta.exposure, beta.outcome, se.exposure, se.outcome, parameters)
      )

      nom <- c(nom, paste0("All - ", subset(mr_method_list(), obj == all_method[i])$name))
    }

    d <- data.frame(
      SNP = c(as.character(x$SNP), nom),
      b = sapply(l, function(y) y$b),
      se = sapply(l, function(y) y$se),
      p = sapply(l, function(y) y$pval),
      samplesize = x$samplesize.outcome[1]
    )
    d$outcome <- x$outcome[1]
    d$exposure <- x$exposure[1]
    return(d)
  })
  res <- subset(
    res,
    select = c(exposure, outcome, id.exposure, id.outcome, samplesize, SNP, b, se, p)
  )
  return(res)
}


#' Forest plot
#'
#' @param singlesnp_results from [mr_singlesnp()].
#' @param exponentiate Plot on exponential scale. The default is `FALSE`.
#'
#' @export
#' @return List of plots
mr_forest_plot <- function(singlesnp_results, exponentiate = FALSE) {
  res <- plyr::dlply(singlesnp_results, c("id.exposure", "id.outcome"), function(d) {
    d <- plyr::mutate(d)
    if (sum(!grepl("All", d$SNP)) < 2) {
      return(
        blank_plot("Insufficient number of SNPs")
      )
    }
    levels(d$SNP)[levels(d$SNP) == "All - Inverse variance weighted"] <- "All - IVW"
    levels(d$SNP)[levels(d$SNP) == "All - MR Egger"] <- "All - Egger"
    am <- grep("All", d$SNP, value = TRUE)
    d$up <- d$b + 1.96 * d$se
    d$lo <- d$b - 1.96 * d$se
    d$tot <- 0.01
    d$tot[d$SNP %in% am] <- 1
    d$SNP <- as.character(d$SNP)
    nom <- d$SNP[!d$SNP %in% am]
    nom <- nom[order(d$b)]
    d <- rbind(d, d[nrow(d), ])
    d$SNP[nrow(d) - 1] <- ""
    d$b[nrow(d) - 1] <- NA
    d$up[nrow(d) - 1] <- NA
    d$lo[nrow(d) - 1] <- NA
    d$SNP <- ordered(d$SNP, levels = c(am, "", nom))

    xint <- 0
    if (exponentiate) {
      d$b <- exp(d$b)
      d$up <- exp(d$up)
      d$lo <- exp(d$lo)
      xint <- 1
    }

    if (utils::packageVersion("ggplot2") <= "3.5.2") {
      ggplot2::ggplot(d, ggplot2::aes(y = SNP, x = b)) +
        ggplot2::geom_vline(xintercept = xint, linetype = "dotted") +
        # ggplot2::geom_errorbarh(ggplot2::aes(xmin=pmax(lo, min(d$b, na.rm=T)), xmax=pmin(up, max(d$b, na.rm=T)), size=as.factor(tot), colour=as.factor(tot)), height=0) +
        ggplot2::geom_errorbarh(
          ggplot2::aes(xmin = lo, xmax = up, linewidth = as.factor(tot), colour = as.factor(tot)),
          height = 0
        ) +
        ggplot2::geom_point(ggplot2::aes(colour = as.factor(tot))) +
        ggplot2::geom_hline(
          ggplot2::aes(yintercept = which(levels(SNP) %in% "")),
          colour = "grey"
        ) +
        ggplot2::scale_colour_manual(values = c("black", "red")) +
        ggplot2::scale_linewidth_manual(values = c(0.3, 1)) +
        # xlim(c(min(c(0, d$b), na.rm=T), max(c(0, d$b), na.rm=T))) +
        ggplot2::theme(
          legend.position = "none",
          axis.text.y = ggplot2::element_text(size = 8),
          axis.ticks.y = ggplot2::element_line(linewidth = 0),
          axis.title.x = ggplot2::element_text(size = 8)
        ) +
        ggplot2::labs(
          y = "",
          x = paste0("MR effect size for\n'", d$exposure[1], "' on '", d$outcome[1], "'")
        )
    } else {
      ggplot2::ggplot(d, ggplot2::aes(y = SNP, x = b)) +
        ggplot2::geom_vline(xintercept = xint, linetype = "dotted") +
        # ggplot2::geom_errorbarh(ggplot2::aes(xmin=pmax(lo, min(d$b, na.rm=T)), xmax=pmin(up, max(d$b, na.rm=T)), size=as.factor(tot), colour=as.factor(tot)), height=0) +
        ggplot2::geom_errorbar(
          ggplot2::aes(xmin = lo, xmax = up, linewidth = as.factor(tot), colour = as.factor(tot)),
          width = 0,
          orientation = "y"
        ) +
        ggplot2::geom_point(ggplot2::aes(colour = as.factor(tot))) +
        ggplot2::geom_hline(
          ggplot2::aes(yintercept = which(levels(SNP) %in% "")),
          colour = "grey"
        ) +
        ggplot2::scale_colour_manual(values = c("black", "red")) +
        ggplot2::scale_linewidth_manual(values = c(0.3, 1)) +
        # xlim(c(min(c(0, d$b), na.rm=T), max(c(0, d$b), na.rm=T))) +
        ggplot2::theme(
          legend.position = "none",
          axis.text.y = ggplot2::element_text(size = 8),
          axis.ticks.y = ggplot2::element_line(linewidth = 0),
          axis.title.x = ggplot2::element_text(size = 8)
        ) +
        ggplot2::labs(
          y = "",
          x = paste0("MR effect size for\n'", d$exposure[1], "' on '", d$outcome[1], "'")
        )
    }
  })
  res
}


#' Density plot
#'
#' @param singlesnp_results from [mr_singlesnp()].
#' @param mr_results Results from [mr()].
#' @param exponentiate Plot on exponentiated scale. The default is `FALSE`.
#' @param bandwidth Density bandwidth parameter.
#'
#' @export
#' @return List of plots
mr_density_plot <- function(
  singlesnp_results,
  mr_results,
  exponentiate = FALSE,
  bandwidth = "nrd0"
) {
  res <- plyr::dlply(singlesnp_results, c("id.exposure", "id.outcome"), function(d) {
    d <- plyr::mutate(d)
    if (sum(!grepl("All", d$SNP)) < 2) {
      return(
        blank_plot("Insufficient number of SNPs")
      )
    }
    d$SNP <- as.character(d$SNP)

    d2 <- subset(d, !grepl("All - ", SNP))
    d1 <- subset(mr_results, id.exposure == d2$id.exposure[1] & id.outcome == d2$id.outcome[1])

    xint <- 0
    if (exponentiate) {
      d$b <- exp(d$b)
      d$up <- exp(d$up)
      d$lo <- exp(d$lo)
      xint <- 1
    }

    ggplot2::ggplot(d2, ggplot2::aes(x = b)) +
      ggplot2::geom_vline(xintercept = xint, linetype = "dotted") +
      ggplot2::geom_density(ggplot2::aes(weight = 1 / se), bw = bandwidth) +
      ggplot2::geom_point(y = 0, colour = "red", ggplot2::aes(size = 1 / se)) +
      ggplot2::geom_vline(data = mr_results, ggplot2::aes(xintercept = b, colour = method)) +
      ggplot2::scale_colour_brewer(type = "qual") +
      ggplot2::labs(x = "Per SNP MR estimate")
  })
  res
}

#' Funnel plot
#'
#' Create funnel plot from single SNP analyses.
#'
#' @param singlesnp_results from [mr_singlesnp()].
#'
#' @export
#' @return List of plots
mr_funnel_plot <- function(singlesnp_results) {
  res <- plyr::dlply(singlesnp_results, c("id.exposure", "id.outcome"), function(d) {
    d <- plyr::mutate(d)
    if (sum(!grepl("All", d$SNP)) < 2) {
      return(
        blank_plot("Insufficient number of SNPs")
      )
    }
    am <- grep("All", d$SNP, value = TRUE)
    d$SNP <- gsub("All - ", "", d$SNP)
    am <- gsub("All - ", "", am)
    ggplot2::ggplot(subset(d, !SNP %in% am), ggplot2::aes(y = 1 / se, x = b)) +
      ggplot2::geom_point() +
      ggplot2::geom_vline(
        data = subset(d, SNP %in% am),
        ggplot2::aes(xintercept = b, colour = SNP)
      ) +
      # ggplot2::scale_colour_brewer(type="qual") +
      ggplot2::scale_colour_manual(
        values = c(
          "#a6cee3",
          "#1f78b4",
          "#b2df8a",
          "#33a02c",
          "#fb9a99",
          "#e31a1c",
          "#fdbf6f",
          "#ff7f00",
          "#cab2d6",
          "#6a3d9a",
          "#ffff99",
          "#b15928"
        )
      ) +
      ggplot2::labs(y = expression(1 / SE[IV]), x = expression(beta[IV]), colour = "MR Method") +
      ggplot2::theme(legend.position = "top", legend.direction = "vertical")
  })
  res
}
